package com.atguigu.bigdata.spark.sql;

import org.apache.spark.SparkConf;
import org.apache.spark.sql.*;
import org.apache.spark.sql.expressions.Aggregator;
/*
         自定义聚合函数类：计算年龄的平均值
     1. 继承org.apache.spark.sql.expressions.Aggregator, 定义泛型
         IN : 输入的数据类型 Long
         BUF : 缓冲区的数据类型 Buff
         OUT : 输出的数据类型 Long
 */
public class Spark03_SparkSQL_UDAF1_JAVA {
    public static void main(String[] args) {
        SparkConf conf = new  SparkConf().setMaster("local[*]").setAppName("sparkSQL");
        SparkSession spark = SparkSession.builder().config(conf).getOrCreate();
        Dataset<Row> df   = spark.read().json("datas/user.json");
        df.createOrReplaceTempView("user");
        spark.udf().register("ageAvg", functions.udaf(new MyAvgUDAF1(), Encoders.LONG()) );
        spark.sql("select ageAvg(age) from user").show();

        spark.stop();

    }
}

class MyAvgUDAF1 extends Aggregator<Long,Buff,Long> {

    // z & zero : 初始值或零值
    // 缓冲区的初始化
    @Override
    public Buff zero() {
        return new Buff(0L,0L);
    }
    // 根据输入的数据更新缓冲区的数据
    @Override
    public Buff reduce(Buff b, Long a) {
        b.setTotal(b.getTotal() + a);
        b.setCount(b.getCount() + 1);
        return b;
    }
    //合并缓冲区
    @Override
    public Buff merge(Buff b1, Buff b2) {
        b1.setCount(b1.getCount() + b2.getCount());
        b1.setTotal(b1.getTotal() + b2.getTotal());
        return b1;
    }

    @Override
    public Long finish(Buff reduction) {
        return reduction.getTotal()/reduction.getCount();
    }
    // 缓冲区的编码操作
    @Override
    public Encoder<Buff> bufferEncoder() {
        return Encoders.bean(Buff.class);
    }
    // 输出的编码操作
    @Override
    public Encoder<Long> outputEncoder() {
        return Encoders.LONG();
    }
}
